package tech.waterism.model;

import cn.hutool.core.io.resource.ResourceUtil;
import tech.waterism.modelbase.PredictBase;
import tech.waterism.modelbase.PredictResult;
import com.alibaba.fastjson.JSON;
import org.apache.log4j.Logger;
import tech.waterism.model.dto.ModelOutputSimple;
import tech.waterism.model.util.ModelOutputHelper;

import java.util.*;

/**
 * 三水源蓄满产流模型（新安江模型）
 * 
 * <p>该模型是新安江模型的核心产流模块，采用蓄满产流机制。
 * 将流域划分为上层、下层和深层三层，分别计算地表径流、壤中流和地下径流。</p>
 * 
 * <p>主要特点：
 * - 三层土壤结构（上层、下层、深层）
 * - 三种径流成分（地表径流RS、壤中流RI、地下径流RG）
 * - 考虑不透水面积产流
 * - 采用蓄满产流和自由水蓄水库分水源方法
 * </p>
 *
 * @author 初京刚
 * @version 2.0
 * @since 2022-01-07
 */
public class SMS_3 extends PredictBase {

    private static final Logger logger = Logger.getLogger(SMS_3.class);

    // ==================== 模型参数 ====================
    
    /** 蓄水容量曲线的方次 */
    private double B;
    
    /** 深层蒸散发系数 */
    private double C;
    
    /** 上层张力水容量占总容量的比例 */
    private double WUMx;
    
    /** 上层张力水容量 (mm) - 计算得到 */
    private double WUM;
    
    /** 下层张力水容量占总容量的比例 */
    private double WLMx;
    
    /** 下层张力水容量 (mm) - 计算得到 */
    private double WLM;
    
    /** 深层张力水容量 (mm) - 计算得到 */
    private double WDM;
    
    /** 流域平均张力水容量 (mm) */
    private double WM;
    
    /** 流域最大张力水容量 (mm) - 计算得到 */
    private double WMM;
    
    /** 地下水消退系数 */
    private double KG;
    
    /** 壤中流消退系数 */
    private double KI;
    
    /** 自由水蓄水库容曲线的方次 */
    private double SMM;
    
    /** 自由水蓄水库容量 (mm) */
    private double SM;
    
    /** 自由水蓄水库容曲线的方次 */
    private double EX;
    
    /** 不透水面积比例 */
    private double IM;
    
    /** 流域蒸发折算系数 */
    private double K;
    
    // ==================== 初始状态 ====================
    
    /** 初始上层张力水含量 (mm) */
    private double WUP;
    
    /** 初始下层张力水含量 (mm) */
    private double WLP;
    
    /** 初始深层张力水含量 (mm) */
    private double WDP;
    
    /** 初始产流面积自由水深 (mm) */
    private double SP;
    
    /** 初始产流面积系数 */
    private double FRP;
    
    // ==================== 输入数据 ====================
    
    /** 流域月潜在蒸散发 PET (mm) - Potential Evapotranspiration */
    private double[] ES;
    
    /** 时段净雨 (mm) */
    private double PE;
    
    /** 降水时间序列 */
    private Date[] TM;
    
    /** 降水量序列 (mm) */
    private double[] PP;
    
    /** 时段长 (h) */
    private double T;
    
    /** 预见期起始时间（用于区分热预期和预见期） */
    private String forecastStartTime;
    
    /** 最后生成的标准化输出对象（用于接口返回） */
    private ModelOutputSimple lastOutput;

    /**
     * 构造函数，从参数Map中初始化模型
     *
     * @param data 模型参数Map
     */
    public SMS_3(Map<String, Object> data) {
        super();
        logger.info("初始化三水源蓄满产流模型...");
        
        try {
            // 提取initialState和parameters子对象（支持新旧两种格式）
            Map<String, Object> initialState = (Map<String, Object>) data.get("initialState");
            Map<String, Object> parameters = (Map<String, Object>) data.get("parameters");
            
            // 如果是新格式（嵌套结构）
            if (initialState != null && parameters != null) {
                logger.info("检测到新格式JSON（嵌套结构）");
                
                // 从initialState读取初始状态
                WDP = parseDouble(initialState, "WDP", "初始深层张力水含量");
                WUP = parseDouble(initialState, "WUP", "初始上层张力水含量");
                WLP = parseDouble(initialState, "WLP", "初始下层张力水含量");
                SP = parseDouble(initialState, "SP", "初始产流面积自由水深");
                FRP = parseDouble(initialState, "FRP", "初始产流面积系数");
                
                // 从parameters读取所有参数
                IM = parseDouble(parameters, "IM", "不透水面积比例");
                B = parseDouble(parameters, "B", "蓄水容量曲线方次");
                K = parseDouble(parameters, "K", "流域蒸发折算系数");
                KG = parseDouble(parameters, "KG", "地下水消退系数");
                KI = parseDouble(parameters, "KI", "壤中流消退系数");
                C = parseDouble(parameters, "C", "深层蒸散发系数");
                WM = parseDouble(parameters, "WM", "流域平均张力水容量");
                SM = parseDouble(parameters, "SM", "自由水蓄水库容量");
                WUMx = parseDouble(parameters, "WUMx", "上层张力水容量比例");
                WLMx = parseDouble(parameters, "WLMx", "下层张力水容量比例");
                EX = parseDouble(parameters, "EX", "自由水蓄水库容曲线方次");
                
                // clen 和 start 从顶层读取（不在parameters中）
                T = parseDouble(data, "clen", "时段长");
                
                // 读取预见期起始时间（可选）
                if (data.containsKey("start")) {
                    forecastStartTime = data.get("start").toString();
                    logger.info("预见期起始时间: " + forecastStartTime);
                }
            } else {
                // 旧格式（扁平化结构）- 向后兼容
                logger.info("检测到旧格式JSON（扁平化结构）");
                
                IM = parseDouble(data, "IM", "不透水面积比例");
                B = parseDouble(data, "B", "蓄水容量曲线方次");
                K = parseDouble(data, "K", "流域蒸发折算系数");
                KG = parseDouble(data, "KG", "地下水消退系数");
                KI = parseDouble(data, "KI", "壤中流消退系数");
                C = parseDouble(data, "C", "深层蒸散发系数");
                WM = parseDouble(data, "WM", "流域平均张力水容量");
                SM = parseDouble(data, "SM", "自由水蓄水库容量");
                WUMx = parseDouble(data, "WUMx", "上层张力水容量比例");
                WLMx = parseDouble(data, "WLMx", "下层张力水容量比例");
                EX = parseDouble(data, "EX", "自由水蓄水库容曲线方次");
                WUP = parseDouble(data, "WUP", "初始上层张力水含量");
                WLP = parseDouble(data, "WLP", "初始下层张力水含量");
                WDP = parseDouble(data, "WDP", "初始深层张力水含量");
                SP = parseDouble(data, "SP", "初始产流面积自由水深");
                FRP = parseDouble(data, "FRP", "初始产流面积系数");
                T = parseDouble(data, "clen", "时段长");
                
                // 读取预见期起始时间（可选）
                if (data.containsKey("start")) {
                    forecastStartTime = data.get("start").toString();
                    logger.info("预见期起始时间: " + forecastStartTime);
                }
            }
            
            // 输入数据 - 解析新的时间序列结构
            parseTimeSeriesData(data);
            
            logger.info("三水源蓄满产流模型初始化完成，时段数: " + PP.length);
        } catch (Exception e) {
            logger.error("三水源蓄满产流模型初始化失败", e);
            throw new RuntimeException("模型初始化失败", e);
        }
    }

    /**
     * 解析参数并记录日志
     */
    private double parseDouble(Map<String, Object> data, String key, String desc) {
        Object value = data.get(key);
        double result;
        if (value instanceof Number) {
            result = ((Number) value).doubleValue();
        } else {
            result = Double.parseDouble(value.toString());
        }
        logger.debug(desc + " (" + key + "): " + result);
        return result;
    }

    /**
     * 解析新的时间序列数据结构
     * 
     * <p>支持以下数据格式：</p>
     * <ul>
     *   <li>新格式：rainfallSeries/timeSeries - 包含时间和降雨量的对象数组</li>
     *   <li>旧格式：dt + rain - 分别存储时间和降雨量的数组</li>
     * </ul>
     * 
     * <p>同时解析月潜在蒸散发数据（PET - Potential Evapotranspiration）：</p>
     * <ul>
     *   <li>monthlyPET - 新字段名（推荐）</li>
     *   <li>monthlyEvaporation - 兼容旧字段名</li>
     *   <li>ES - 最旧格式（数组形式）</li>
     * </ul>
     * 
     * @param data JSON数据Map，包含降雨时间序列和蒸散发数据
     * @throws RuntimeException 当无法解析降雨数据时抛出
     */
    private void parseTimeSeriesData(Map<String, Object> data) {
        logger.info("解析时间序列数据...");
        
        // 1. 解析主降雨序列 (rainfallSeries，兼容旧名称timeSeries)
        String seriesKey = data.containsKey("rainfallSeries") ? "rainfallSeries" : "timeSeries";
        if (data.containsKey(seriesKey)) {
            List<Map<String, Object>> timeSeriesList = (List<Map<String, Object>>) data.get(seriesKey);
            int length = timeSeriesList.size();
            
            TM = new Date[length];
            PP = new double[length];
            
            for (int i = 0; i < length; i++) {
                Map<String, Object> point = timeSeriesList.get(i);
                
                // 解析时间
                String timeStr = (String) point.get("time");
                try {
                    TM[i] = new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss").parse(timeStr);
                } catch (Exception e) {
                    logger.error("时间解析失败: " + timeStr, e);
                    throw new RuntimeException("时间解析失败", e);
                }
                
                // 解析降雨量
                Object rainObj = point.get("rain");
                PP[i] = rainObj instanceof Number ? ((Number) rainObj).doubleValue() : Double.parseDouble(rainObj.toString());
            }
            
            logger.info("主降雨序列解析完成: " + length + " 个时段");
        } else {
            // 兼容旧格式
            logger.warn("未找到rainfallSeries/timeSeries，尝试解析旧格式...");
            TM = parseArray(data, "dt", Date[].class, "降水时间");
            PP = parseArray(data, "rain", double[].class, "降水量");
            
            // 检查解析结果
            if (PP == null || TM == null) {
                throw new RuntimeException("无法解析降雨数据：未找到rainfallSeries/timeSeries或旧格式的dt/rain字段");
            }
        }
        
        // 2. 解析月潜在蒸散发 (monthlyPET - Potential Evapotranspiration)
        // 优先从parameters中查找（新格式）
        Map<String, Object> parameters = (Map<String, Object>) data.get("parameters");
        Object monthlyPETSource = parameters != null ? parameters.get("monthlyPET") : data.get("monthlyPET");
        
        if (monthlyPETSource != null) {
            List<Map<String, Object>> monthlyList = (List<Map<String, Object>>) monthlyPETSource;
            ES = new double[12];
            
            for (int i = 0; i < Math.min(12, monthlyList.size()); i++) {
                Map<String, Object> monthData = monthlyList.get(i);
                Object petObj = monthData.get("pet");
                ES[i] = petObj instanceof Number ? ((Number) petObj).doubleValue() : Double.parseDouble(petObj.toString());
            }
            
            logger.info("月潜在蒸散发(PET)解析完成: " + ES.length + " 个月");
        } else if (data.containsKey("monthlyEvaporation")) {
            // 兼容旧格式 monthlyEvaporation
            logger.warn("使用旧字段名 monthlyEvaporation，建议使用 monthlyPET");
            List<Map<String, Object>> monthlyList = (List<Map<String, Object>>) data.get("monthlyEvaporation");
            ES = new double[12];
            
            for (int i = 0; i < Math.min(12, monthlyList.size()); i++) {
                Map<String, Object> monthData = monthlyList.get(i);
                Object evapObj = monthData.get("evaporation");
                if (evapObj == null) {
                    evapObj = monthData.get("pet");
                }
                ES[i] = evapObj instanceof Number ? ((Number) evapObj).doubleValue() : Double.parseDouble(evapObj.toString());
            }
            
            logger.info("月潜在蒸散发(PET)解析完成: " + ES.length + " 个月");
        } else if (data.containsKey("ES")) {
            // 兼容最旧格式
            logger.warn("未找到monthlyPET，使用旧格式ES");
            ES = parseArray(data, "ES", double[].class, "月潜在蒸散发");
        }
        
        logger.info("时间序列数据解析完成");
    }

    /**
     * 解析数组字段 - 兼容字符串和直接的JSON数组
     */
    private <T> T parseArray(Map<String, Object> data, String key, Class<T> clazz, String desc) {
        Object value = data.get(key);
        T result;
        if (value instanceof String) {
            // 旧格式：字符串形式的JSON数组
            result = JSON.parseObject((String) value, clazz);
        } else {
            // 新格式：已经是JSON对象
            result = JSON.parseObject(JSON.toJSONString(value), clazz);
        }
        logger.debug(desc + " (" + key + "): 已解析，长度=" + 
                    (result != null && result.getClass().isArray() ? 
                     java.lang.reflect.Array.getLength(result) : "unknown"));
        return result;
    }

    /**
     * 测试入口 - 使用转换后的标准JSON文件
     */
    public static void main(String[] args) {
        logger.info("开始三水源蓄满产流模型测试...");
        logger.info("读取转换后的JSON文件: json/SMS_3-input.json");
        
        // 读取转换后的标准JSON文件
        String content = ResourceUtil.readUtf8Str("json/SMS_3-input.json");
        Map<String, Object> data = JSON.parseObject(content, Map.class);
        
        SMS_3 model = new SMS_3(data);
        PredictResult result = model.predict();
        
        logger.info("模型计算完成");
        logger.info("输出结果: " + JSON.toJSONString(result));
        System.out.println(JSON.toJSONString(result));
    }

    @Override
    public void init() {
        // 初始化已在构造函数中完成
    }

    @Override
    public PredictResult predict() {
        logger.info("开始三水源蓄满产流预测计算...");
        
        PredictResult result = new PredictResult();
        Map<String, double[]> map = runModel();
        
        // ===== 传统数组格式输出（向后兼容）=====
        // 分水源产流结果
        double[][] RSim = new double[3][];
        RSim[0] = map.get("RS");  // 流域地表产流
        RSim[1] = map.get("RI");  // 流域土壤产流（壤中流）
        RSim[2] = map.get("RG");  // 流域地下产流
        
        // 流域面积总产流
        result.setRunoffSim(map.get("RunoffSim"));
        result.setRSim(RSim);
        
        // 时间序列
        result.setDt(TM);
        result.setRain(PP);
        
        // 更新状态
        Map<String, List<String>> status = new HashMap<>();
        status.put("WUP", toStringList(map.get("WU")));
        status.put("WLP", toStringList(map.get("WL")));
        status.put("WDP", toStringList(map.get("WD")));
        status.put("SP", toStringList(map.get("S")));
        status.put("FRP", toStringList(map.get("FR")));
        result.setNewStatus(status);
        
        // ===== 新增：扁平化输出并保存到文件 =====
        try {
            ModelOutputHelper helper = new ModelOutputHelper("SMS_3", "1.0");
            
            // 计算预热期时段数
            int warmUpPeriods = calculateWarmUpPeriods();
            
            // 设置时间序列和预热期
            helper.setTimeStamps(TM)
                  .setWarmUpPeriods(warmUpPeriods);
            
            // 添加所有水文要素（扁平化数组）
            helper.addDataArray("rain", PP)
                  .addDataArray("RunoffSim", map.get("RunoffSim"))
                  .addDataArray("RS", map.get("RS"))
                  .addDataArray("RI", map.get("RI"))
                  .addDataArray("RG", map.get("RG"));
            
            // 添加状态变量时间序列
            double[] WU = map.get("WU");
            double[] WL = map.get("WL");
            double[] WD = map.get("WD");
            double[] S = map.get("S");
            double[] FR = map.get("FR");
            
            helper.addStateArray("WUP", WU)
                  .addStateArray("WLP", WL)
                  .addStateArray("WDP", WD)
                  .addStateArray("SP", S)
                  .addStateArray("FRP", FR);
            
            // 设置最终状态
            Map<String, Object> finalState = new HashMap<>();
            finalState.put("WUP", WU[WU.length - 1]);
            finalState.put("WLP", WL[WL.length - 1]);
            finalState.put("WDP", WD[WD.length - 1]);
            finalState.put("SP", S[S.length - 1]);
            finalState.put("FRP", FR[FR.length - 1]);
            helper.setFinalState(finalState);
            
            // 设置模型参数
            Map<String, Double> params = new HashMap<>();
            params.put("WM_流域平均张力水容量", WM);
            params.put("B_蓄水容量曲线指数", B);
            params.put("SM_自由水蓄水容量", SM);
            params.put("KG_地下径流消退系数", KG);
            params.put("KI_壤中流消退系数", KI);
            helper.setModelParameters(params);
            
            // 设置元数据
            Map<String, Object> metadata = new HashMap<>();
            metadata.put("timestep_hours", T);
            helper.setMetadata(metadata);
            
            // 构建并保存到文件
            ModelOutputSimple output = helper.build();
            this.lastOutput = output;  // 保存输出对象供接口使用
            String outputPath = "output/SMS_3_output.json";
            ModelOutputHelper.saveToFile(output, outputPath);
            logger.info("✅ SMS_3模型输出已保存到: " + outputPath);
        } catch (Exception e) {
            logger.error("保存模型输出时出错: " + e.getMessage(), e);
        }
        
        logger.info("三水源蓄满产流预测计算完成");
        return result;
    }
    
    /**
     * 获取最后生成的标准化输出对象
     * 
     * <p>该方法用于接口返回，提供与output文件一致的格式</p>
     * 
     * @return ModelOutputSimple对象，包含完整的模型输出信息
     */
    public ModelOutputSimple getLastOutput() {
        return lastOutput;
    }
    
    /**
     * 计算预热期时段数
     * 
     * <p>根据预见期起始时间（start）在时间序列中查找对应的位置，
     * 该位置之前的时段作为预热期（warm-up period），之后的作为预见期（forecast period）</p>
     * 
     * @return 预热期时段数，如果未指定start则返回默认值（总时段数>72返回72，否则返回0）
     */
    private int calculateWarmUpPeriods() {
        if (forecastStartTime == null || forecastStartTime.isEmpty()) {
            // 如果没有指定start，使用默认策略
            int defaultWarmUp = PP.length > 72 ? 72 : 0;
            logger.info("未指定预见期起始时间，使用默认预热期: " + defaultWarmUp + " 时段");
            return defaultWarmUp;
        }
        
        try {
            // 解析预见期起始时间
            java.text.SimpleDateFormat sdf = new java.text.SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
            Date startDate = sdf.parse(forecastStartTime);
            
            // 在时间序列中查找对应位置
            for (int i = 0; i < TM.length; i++) {
                if (TM[i].equals(startDate) || TM[i].after(startDate)) {
                    logger.info("预见期起始于第 " + i + " 时段（时间: " + sdf.format(TM[i]) + "），预热期: " + i + " 时段");
                    return i;
                }
            }
            
            // 如果start晚于所有时间，则全部作为预热期
            logger.warn("预见期起始时间 " + forecastStartTime + " 晚于所有时段，全部视为预热期");
            return PP.length;
            
        } catch (Exception e) {
            logger.error("解析预见期起始时间失败: " + forecastStartTime + "，使用默认策略", e);
            return PP.length > 72 ? 72 : 0;
        }
    }

    /**
     * 运行三水源蓄满产流模型主流程
     * 
     * <p>新安江模型核心产流算法，采用三层土壤结构：</p>
     * <ul>
     *   <li>上层（Upper Layer）- WUM：快速响应降雨</li>
     *   <li>下层（Lower Layer）- WLM：中等响应</li>
     *   <li>深层（Deep Layer）- WDM：慢速响应</li>
     * </ul>
     * 
     * <p>产流机制：</p>
     * <pre>
     * 降雨 → 蒸散发 → 净雨 → 张力水蓄水 → 自由水蓄水 → 三水源径流分离
     *                         ↓
     *                    上、下、深三层
     * </pre>
     * 
     * <p>计算流程：</p>
     * <ol>
     *   <li>计算各层张力水容量（WUM, WLM, WDM）</li>
     *   <li>计算分层参数（WMM, SMM）</li>
     *   <li>调整消退系数（确保稳定性）</li>
     *   <li>逐时段计算产流</li>
     * </ol>
     * 
     * <p>产流分量：</p>
     * <ul>
     *   <li>RS - 地表径流（Surface Runoff）：快速响应，峰值大</li>
     *   <li>RI - 壤中流（Interflow）：中等响应</li>
     *   <li>RG - 地下径流（Groundwater）：慢速响应，基流</li>
     * </ul>
     *
     * @return 模拟结果Map，包含以下键值：
     *         <ul>
     *           <li>RunoffSim - 总产流序列 (mm)</li>
     *           <li>RS - 地表径流序列 (mm)</li>
     *           <li>RI - 壤中流序列 (mm)</li>
     *           <li>RG - 地下径流序列 (mm)</li>
     *           <li>WU - 上层张力水时间序列 (mm)</li>
     *           <li>WL - 下层张力水时间序列 (mm)</li>
     *           <li>WD - 深层张力水时间序列 (mm)</li>
     *           <li>S - 自由水蓄量时间序列 (mm)</li>
     *           <li>FR - 产流面积系数时间序列</li>
     *         </ul>
     */
    public Map<String, double[]> runModel() {
        logger.info("开始三水源蓄满产流计算...");
        
        // 计算各层张力水容量
        WUM = WUMx * WM;
        WLM = (1 - WUMx) * WLMx * WM;
        WDM = WM - WUM - WLM;
        double WP = WUP + WLP + WDP;
        
        logger.debug("上层容量: " + WUM + ", 下层容量: " + WLM + ", 深层容量: " + WDM);
        
        // 分层计算参数
        double DIV = 5.0;  // 分层计算阈值
        WMM = (1.0 + B) * WM / (1.0 - IM);
        SMM = (1.0 + EX) * SM;
        
        int Pnum = PP.length;
        
        // 初始化状态数组
        double[] WU = new double[Pnum + 1];  // 上层张力水
        double[] WL = new double[Pnum + 1];  // 下层张力水
        double[] WD = new double[Pnum + 1];  // 深层张力水
        double[] W = new double[Pnum + 1];   // 总张力水
        double[] S = new double[Pnum + 1];   // 自由水蓄量
        double[] FR = new double[Pnum + 1];  // 产流面积系数
        
        // 产流结果数组
        double[] RSP = new double[Pnum];  // 地表径流
        double[] RIP = new double[Pnum];  // 壤中流
        double[] RGP = new double[Pnum];  // 地下径流
        double[] RP = new double[Pnum];   // 总径流
        
        // 设置初始状态
        WU[0] = WUP;
        WL[0] = WLP;
        WD[0] = WDP;
        W[0] = WP;
        S[0] = SP;
        FR[0] = FRP;
        
        // 调整消退系数确保稳定性
        if (KG + KI > 0.9) {
            double tmp = (KG + KI - 0.9) / (KG + KI);
            KG = KG - KG * tmp;
            KI = KI - KI * tmp;
            logger.warn("消退系数过大，已调整 KG: " + KG + ", KI: " + KI);
        }
        
        // 时段消退系数转换
        double HGI = (1.0 - Math.pow(1.0 - KG - KI, T / 24.0)) / (KG + KI);
        KG = HGI * KG;
        KI = HGI * KI;
        
        // 逐时段计算
        for (int i = 0; i < Pnum; i++) {
            calculateTimeStep(i, WU, WL, WD, W, S, FR, RSP, RIP, RGP, RP, DIV);
        }
        
        // 组装结果
        Map<String, double[]> result = new HashMap<>();
        result.put("RunoffSim", RP);
        result.put("RS", RSP);
        result.put("RI", RIP);
        result.put("RG", RGP);
        result.put("WU", WU);
        result.put("WL", WL);
        result.put("WD", WD);
        result.put("S", S);
        result.put("FR", FR);
        
        logger.debug("三水源蓄满产流计算完成");
        return result;
    }

    /**
     * 单时段产流计算
     * 
     * <p>计算步骤：</p>
     * <ol>
     *   <li>计算时段潜在蒸散发（PET）</li>
     *   <li>计算净雨 PE = 降雨 - K×PET</li>
     *   <li>根据净雨大小选择计算方法：</li>
     *   <ul>
     *     <li>PE ≥ 2×DIV：大雨，分层计算（提高精度）</li>
     *     <li>PE < 2×DIV：小雨，不分层</li>
     *   </ul>
     *   <li>确保产流非负</li>
     *   <li>计算总产流 RP = RS + RI + RG</li>
     * </ol>
     * 
     * <p>分层计算原理：</p>
     * <p>大雨时，将净雨分成多个小层次（ND层）逐层计算，
     * 可以更准确地模拟蓄水容量曲线的非线性特性</p>
     *
     * @param i   时段索引
     * @param WU  上层张力水数组 (mm)
     * @param WL  下层张力水数组 (mm)
     * @param WD  深层张力水数组 (mm)
     * @param W   总张力水数组 (mm)
     * @param S   自由水蓄量数组 (mm)
     * @param FR  产流面积系数数组
     * @param RSP 地表径流数组 (mm)
     * @param RIP 壤中流数组 (mm)
     * @param RGP 地下径流数组 (mm)
     * @param RP  总产流数组 (mm)
     * @param DIV 分层计算阈值 (mm)
     */
    private void calculateTimeStep(int i, double[] WU, double[] WL, double[] WD, double[] W,
                                   double[] S, double[] FR, double[] RSP, double[] RIP, 
                                   double[] RGP, double[] RP, double DIV) {
        // 计算月均天数
        int IDAY = getDaysInMonth(TM[i]);
        
        // 计算时段潜在蒸散发量 (PET)
        Calendar time = Calendar.getInstance();
        time.setTime(TM[i]);
        int Mth = time.get(Calendar.MONTH) + 1;
        double EM = ES[Mth - 1] / (IDAY * 24.0 / T);  // 月PET转换为时段PET
        double EK = K * EM;  // K为流域蒸散发折算系数
        
        // 计算净雨
        PE = PP[i] - EK;
        
        // 根据净雨大小选择计算方法
        if (PE >= 2 * DIV) {
            // 大雨，分层计算
            calculateWithLayering(i, WU, WL, WD, W, S, FR, RSP, RIP, RGP, EK, DIV);
        } else {
            // 小雨，不分层
            calculateWithoutLayering(i, WU, WL, WD, W, S, FR, RSP, RIP, RGP, EK);
        }
        
        // 确保产流非负
        RSP[i] = Math.max(0.0, RSP[i]);
        RIP[i] = Math.max(0.0, RIP[i]);
        RGP[i] = Math.max(0.0, RGP[i]);
        
        // 计算总产流
        RP[i] = RSP[i] + RIP[i] + RGP[i];
        
        if (i % 100 == 0 || RP[i] > 0) {
            logger.debug("时段 " + i + ", 降雨: " + PP[i] + ", 产流: " + RP[i] 
                    + " (RS: " + RSP[i] + ", RI: " + RIP[i] + ", RG: " + RGP[i] + ")");
        }
    }

    /**
     * 获取月份天数（简化计算）
     * 
     * <p>返回值：</p>
     * <ul>
     *   <li>1, 3, 5, 7, 8, 10, 12月 → 31天</li>
     *   <li>4, 6, 9, 11月 → 30天</li>
     *   <li>2月 → 28天（未考虑闰年）</li>
     * </ul>
     * 
     * <p>注意：此方法未考虑闰年，2月始终返回28天</p>
     *
     * @param date 日期
     * @return 该月天数
     */
    private int getDaysInMonth(Date date) {
        Calendar time = Calendar.getInstance();
        time.setTime(date);
        int month = time.get(Calendar.MONTH) + 1;
        
        switch (month) {
            case 4:
            case 6:
            case 9:
            case 11:
                return 30;
            case 2:
                return 28;
            default:
                return 31;
        }
    }

    /**
     * 大雨分层计算（提高精度）
     * 
     * <p>分层原理：</p>
     * <p>将净雨PE分成ND个小层（每层约DIV mm），逐层计算产流，
     * 可以更好地模拟蓄水容量曲线的非线性特性，避免大雨时精度损失</p>
     * 
     * <p>计算步骤：</p>
     * <ol>
     *   <li>计算分层数 ND = floor(PE / DIV)</li>
     *   <li>将PE分配到各层（前ND-1层为DIV，最后一层为余数）</li>
     *   <li>调用张力水蓄水量计算</li>
     *   <li>调用自由水蓄水库及分水源计算</li>
     * </ol>
     * 
     * <p>优点：</p>
     * <ul>
     *   <li>避免大雨时一次性计算导致的误差累积</li>
     *   <li>更准确地反映蓄满产流的渐进过程</li>
     * </ul>
     *
     * @param i    时段索引
     * @param WU   上层张力水数组
     * @param WL   下层张力水数组
     * @param WD   深层张力水数组
     * @param W    总张力水数组
     * @param S    自由水蓄量数组
     * @param FR   产流面积系数数组
     * @param RSP  地表径流数组
     * @param RIP  壤中流数组
     * @param RGP  地下径流数组
     * @param EK   时段实际蒸散发量 (mm)
     * @param DIV  分层阈值 (mm)
     */
    private void calculateWithLayering(int i, double[] WU, double[] WL, double[] WD, double[] W,
                                      double[] S, double[] FR, double[] RSP, double[] RIP, 
                                      double[] RGP, double EK, double DIV) {
        int ND = (int) Math.floor(PE / DIV);
        double[] PED = new double[ND];
        double[] RD = new double[ND];
        
        for (int II = 0; II < ND - 1; II++) {
            PED[II] = DIV;
        }
        PED[ND - 1] = PE - (ND - 1) * DIV;
        
        // 张力水蓄水量计算
        Map<String, Double> WMap = calculateTensionWaterStorage(
                EK, W[i], WU[i], WL[i], WD[i], RD, PED, ND);
        W[i + 1] = WMap.get("W");
        WU[i + 1] = WMap.get("WU");
        WL[i + 1] = WMap.get("WL");
        WD[i + 1] = WMap.get("WD");
        
        // 自由水蓄水量及分水源计算
        Map<String, Double> RMap = calculateFreeWaterAndRunoff(
                FR[i], S[i], RD, PED, ND);
        S[i + 1] = RMap.get("S");
        FR[i + 1] = RMap.get("FR");
        RSP[i] = RMap.get("RS");
        RIP[i] = RMap.get("RI");
        RGP[i] = RMap.get("RG");
    }

    /**
     * 小雨不分层计算（快速计算）
     * 
     * <p>当净雨PE较小时（PE < 2×DIV），直接一次性计算，无需分层</p>
     * 
     * <p>计算步骤：</p>
     * <ol>
     *   <li>设置层数 ND = 1</li>
     *   <li>调用张力水蓄水量计算</li>
     *   <li>调用自由水蓄水库及分水源计算</li>
     * </ol>
     * 
     * <p>适用场景：</p>
     * <ul>
     *   <li>小雨或中雨</li>
     *   <li>计算速度优先，精度要求不高的情况</li>
     * </ul>
     *
     * @param i    时段索引
     * @param WU   上层张力水数组
     * @param WL   下层张力水数组
     * @param WD   深层张力水数组
     * @param W    总张力水数组
     * @param S    自由水蓄量数组
     * @param FR   产流面积系数数组
     * @param RSP  地表径流数组
     * @param RIP  壤中流数组
     * @param RGP  地下径流数组
     * @param EK   时段实际蒸散发量 (mm)
     */
    private void calculateWithoutLayering(int i, double[] WU, double[] WL, double[] WD, double[] W,
                                         double[] S, double[] FR, double[] RSP, double[] RIP, 
                                         double[] RGP, double EK) {
        int ND = 1;
        double[] PED = new double[]{PE};
        double[] RD = new double[ND];
        
        Map<String, Double> WMap = calculateTensionWaterStorage(
                EK, W[i], WU[i], WL[i], WD[i], RD, PED, ND);
        W[i + 1] = WMap.get("W");
        WU[i + 1] = WMap.get("WU");
        WL[i + 1] = WMap.get("WL");
        WD[i + 1] = WMap.get("WD");
        
        Map<String, Double> RMap = calculateFreeWaterAndRunoff(
                FR[i], S[i], RD, PED, ND);
        S[i + 1] = RMap.get("S");
        FR[i + 1] = RMap.get("FR");
        RSP[i] = RMap.get("RS");
        RIP[i] = RMap.get("RI");
        RGP[i] = RMap.get("RG");
    }

    /**
     * 张力水蓄水量计算 - 新安江模型核心算法之一
     * 
     * <p>张力水概念：</p>
     * <p>张力水是土壤中被毛管力吸附的水分，不能自由流动产流。
     * 只有当张力水接近饱和时，多余的水分才能转化为自由水产流。</p>
     * 
     * <p>蓄满产流机制：</p>
     * <pre>
     * 降雨 → 优先补充张力水 → 张力水饱和 → 产生自由水 → 径流
     * </pre>
     * 
     * <p>蓄水容量分布曲线（抛物线）：</p>
     * <pre>
     * A = WMM × [1 - (1 - W/WM)^(1/(1+B))]
     * </pre>
     * 其中：
     * <ul>
     *   <li>A - 点蓄水容量</li>
     *   <li>WMM - 最大张力水容量 = (1+B)×WM/(1-IM)</li>
     *   <li>B - 蓄水容量曲线指数</li>
     * </ul>
     * 
     * <p>计算内容：</p>
     * <ol>
     *   <li>蒸发情况（PE ≤ 0）：</li>
     *   <ul>
     *     <li>优先从上层蓄水蒸发</li>
     *     <li>上层不足时，从下层蒸发（按比例）</li>
     *     <li>下层不足时，从深层蒸发</li>
     *   </ul>
     *   <li>产流情况（PE > 0）：</li>
     *   <ul>
     *     <li>根据蓄水容量曲线计算产流量RD[]</li>
     *     <li>将入渗水量分配到上、下、深三层</li>
     *     <li>各层优先级：上层 → 下层 → 深层</li>
     *   </ul>
     * </ol>
     * 
     * <p>参数说明：</p>
     * <ul>
     *   <li>WUM - 上层张力水容量 (mm)</li>
     *   <li>WLM - 下层张力水容量 (mm)</li>
     *   <li>WDM - 深层张力水容量 (mm)</li>
     *   <li>C - 深层蒸散发系数</li>
     * </ul>
     *
     * @param EK  时段蒸散发能力 (mm)
     * @param W   当前总张力水含量 (mm)
     * @param WU  当前上层张力水含量 (mm)
     * @param WL  当前下层张力水含量 (mm)
     * @param WD  当前深层张力水含量 (mm)
     * @param RD  各层产流量数组（输出参数）(mm)
     * @param PED 各层净雨量数组 (mm)
     * @param ND  分层数
     * @return 各层张力水含量Map，包含：W, WU, WL, WD
     */
    private Map<String, Double> calculateTensionWaterStorage(
            double EK, double W, double WU, double WL, double WD,
            double[] RD, double[] PED, int ND) {
        
        if (PE <= 0.0) {
            // 蒸发情况
            if (WU + PE >= 0.0) {
                WU = WU + PE;
            } else {
                double EU = WU + EK + PE;
                WU = 0.0;
                double EL = (EK - EU) * WL / WLM;
                
                if (WL < C * WLM) {
                    EL = C * (EK - EU);
                }
                
                if (WL - EL < 0.0) {
                    double ED = EL - WL;
                    WL = 0.0;
                    WD = WD - ED;
                } else {
                    WL = WL - EL;
                }
            }
            W = WU + WL + WD;
        } else {
            // 产流情况
            double A;
            if (WM - W < 0.0001) {
                A = WMM;
            } else {
                A = WMM * (1.0 - Math.pow(1.0 - W / WM, 1.0 / (1.0 + B)));
            }
            
            double R = 0.0;
            double PEDS = 0.0;
            
            for (int I = 0; I < ND; I++) {
                A = A + PED[I];
                PEDS = PEDS + PED[I];
                double RI = R;
                R = PEDS - WM + W;
                
                if (A < WMM) {
                    R = R + WM * Math.pow(1.0 - A / WMM, 1.0 + B);
                }
                
                RD[I] = R - RI;
            }
            
            // 分配到各层
            if (WU + PE - R <= WUM) {
                WU = WU + PE - R;
            } else {
                if (WU + WL + PE - R - WUM >= WLM) {
                    WU = WUM;
                    WL = WLM;
                    WD = W + PEDS - R - WU - WL;
                    if (WD > WDM) {
                        WD = WDM;
                    }
                } else {
                    WL = WU + WL + PE - R - WUM;
                    WU = WUM;
                }
            }
            
            W = WU + WL + WD;
        }
        
        Map<String, Double> result = new HashMap<>();
        result.put("W", W);
        result.put("WU", WU);
        result.put("WL", WL);
        result.put("WD", WD);
        return result;
    }

    /**
     * 自由水蓄水库及分水源计算 - 新安江模型核心算法之二
     * 
     * <p>自由水概念：</p>
     * <p>自由水是土壤中可以自由流动的水分，当张力水饱和后，
     * 多余的水分转化为自由水，在自由水蓄水库中进行分配。</p>
     * 
     * <p>三水源分离机制：</p>
     * <pre>
     * 自由水 → 地表径流（RS）：快速响应，流速快
     *        → 壤中流（RI）：中等响应，在土壤中侧向流动
     *        → 地下径流（RG）：慢速响应，地下水补给
     * </pre>
     * 
     * <p>自由水蓄水库容曲线（抛物线）：</p>
     * <pre>
     * AU = SMM × [1 - (1 - S/SM)^(1/(1+EX))]
     * </pre>
     * 其中：
     * <ul>
     *   <li>AU - 点自由水容量</li>
     *   <li>SMM - 最大自由水容量 = (1+EX)×SM</li>
     *   <li>EX - 自由水蓄水库容曲线指数</li>
     * </ul>
     * 
     * <p>产流面积概念：</p>
     * <ul>
     *   <li>FR - 产流面积系数（0-1）</li>
     *   <li>产流量 = 单位面积产流深度 × FR</li>
     *   <li>FR动态变化，随降雨增大而增大</li>
     * </ul>
     * 
     * <p>计算流程：</p>
     * <ol>
     *   <li>无降雨情况（PE ≤ 0）：</li>
     *   <ul>
     *     <li>RG = S × KG × FR（地下水退水）</li>
     *     <li>RI = S × KI × FR（壤中流退水）</li>
     *     <li>RS = 0</li>
     *     <li>S = S × (1 - KG - KI)（自由水消退）</li>
     *   </ul>
     *   <li>有降雨情况（PE > 0）：</li>
     *   <ul>
     *     <li>计算不透水面积产流 RB = IM × PE</li>
     *     <li>更新产流面积系数 FR</li>
     *     <li>根据自由水蓄水库容曲线计算地表径流 RR</li>
     *     <li>从自由水中分离出壤中流和地下径流：</li>
     *     <ul>
     *       <li>RG = S × KGD × FR</li>
     *       <li>RI = S × KID × FR</li>
     *     </ul>
     *     <li>RS = RR + RB（总地表径流）</li>
     *   </ul>
     * </ol>
     * 
     * <p>消退系数：</p>
     * <ul>
     *   <li>KG - 地下水消退系数（日值）</li>
     *   <li>KI - 壤中流消退系数（日值）</li>
     *   <li>KGD, KID - 时段消退系数，由日值换算</li>
     * </ul>
     * 
     * <p>参数说明：</p>
     * <ul>
     *   <li>SM - 自由水蓄水库容量 (mm)</li>
     *   <li>EX - 自由水蓄水库容曲线指数</li>
     *   <li>IM - 不透水面积比例</li>
     *   <li>KG - 地下水消退系数</li>
     *   <li>KI - 壤中流消退系数</li>
     * </ul>
     *
     * @param FR  当前产流面积系数
     * @param S   当前自由水蓄量 (mm)
     * @param RD  各层产流量数组 (mm)
     * @param PED 各层净雨量数组 (mm)
     * @param ND  分层数
     * @return 径流分量Map，包含：FR, S, RS, RI, RG
     */
    private Map<String, Double> calculateFreeWaterAndRunoff(
            double FR, double S, double[] RD, double[] PED, int ND) {
        
        if (PE <= 0.0) {
            // 无降雨，退水过程
            double RS = 0.0;
            double RG = S * KG * FR;
            double RI = S * KI * FR;
            S = S * (1.0 - KG - KI);
            
            Map<String, Double> result = new HashMap<>();
            result.put("FR", FR);
            result.put("S", S);
            result.put("RS", RS);
            result.put("RI", RI);
            result.put("RG", RG);
            return result;
        } else {
            // 有降雨，产流过程
            double RB = IM * PE;  // 不透水面积产流
            
            double KID = (1.0 - Math.pow(1.0 - (KG + KI), 1.0 / ND)) / (KG + KI);
            double KGD = KID * KG;
            KID = KID * KI;
            
            double RS = 0.0;
            double RI = 0.0;
            double RG = 0.0;
            
            for (int I = 0; I < ND; I++) {
                double TD = RD[I] - IM * PED[I];
                double X = FR;
                FR = TD / PED[I];
                S = X * S / FR;
                
                double RR;
                if (S >= SM) {
                    RR = (PED[I] + S - SM) * FR;
                } else {
                    double AU = SMM * (1.0 - Math.pow(1.0 - S / SM, 1.0 / (1.0 + EX)));
                    if (AU + PED[I] < SMM) {
                        RR = (PED[I] - SM + S + SM * Math.pow(1.0 - (PED[I] + AU) / SMM, 1.0 + EX)) * FR;
                    } else {
                        RR = (PED[I] + S - SM) * FR;
                    }
                }
                
                RS = RR + RS;
                S = PED[I] - RR / FR + S;
                RG = S * KGD * FR + RG;
                RI = S * KID * FR + RI;
                S = S * (1.0 - KID - KGD);
            }
            
            RS = RS + RB;  // 加上不透水面积产流
            
            Map<String, Double> result = new HashMap<>();
            result.put("FR", FR);
            result.put("S", S);
            result.put("RS", RS);
            result.put("RI", RI);
            result.put("RG", RG);
            return result;
        }
    }

    /**
     * double数组转String列表
     */
    private List<String> toStringList(double[] array) {
        List<String> list = new ArrayList<>(array.length);
        for (double v : array) {
            list.add(String.valueOf(v));
        }
        return list;
    }
}

